import matplotlib.pyplot as plt
import numpy as np
import sys
import os

# 添加src目录到路径
sys.path.append('src')

import util
from linear_model import LinearModel
from p01b_logreg import LogisticRegression
from p01e_gda import GDA

# 数据集路径（从src目录运行时需要回到上级目录）
ds1_training_set_path = '../data/ds1_train.csv'
ds1_valid_set_path = '../data/ds1_valid.csv'
ds2_training_set_path = '../data/ds2_train.csv'
ds2_valid_set_path = '../data/ds2_valid.csv'


def plot_comparison(x, y, theta_lr, theta_gda, title, save_path=None):
    """使用util.py的plot函数风格，但支持两个决策边界"""
    # 使用util.py的plot函数风格
    plt.figure(figsize=(10, 8))

    # 绘制数据集（与util.py中的plot函数相同）
    plt.plot(x[y == 1, -2], x[y == 1, -1], 'bx', linewidth=2, markersize=8, label='y=1')
    plt.plot(x[y == 0, -2], x[y == 0, -1], 'go', linewidth=2, markersize=8, label='y=0')

    # 绘制决策边界（基于util.py中的逻辑）
    margin1 = (max(x[:, -2]) - min(x[:, -2])) * 0.2
    margin2 = (max(x[:, -1]) - min(x[:, -1])) * 0.2
    x1 = np.arange(min(x[:, -2]) - margin1, max(x[:, -2]) + margin1, 0.01)

    # 逻辑回归决策边界（红色实线）
    x2_lr = -(theta_lr[0] / theta_lr[2] + theta_lr[1] / theta_lr[2] * x1)
    plt.plot(x1, x2_lr, c='red', linewidth=3, label='Logistic Regression')

    # GDA决策边界（黑色虚线）
    x2_gda = -(theta_gda[0] / theta_gda[2] + theta_gda[1] / theta_gda[2] * x1)
    plt.plot(x1, x2_gda, c='black', linewidth=3, label='GDA', linestyle='--')

    # 设置图形属性（与util.py中的plot函数相同）
    plt.xlim(x[:, -2].min() - margin1, x[:, -2].max() + margin1)
    plt.ylim(x[:, -1].min() - margin2, x[:, -1].max() + margin2)
    plt.xlabel('x1', fontsize=12)
    plt.ylabel('x2', fontsize=12)
    plt.title(title, fontsize=14)
    plt.legend(loc="upper left", fontsize=10)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()

    if save_path is not None:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')

    plt.show()


def main():
    # ===== 数据集1分析 =====
    print("\n=== 数据集1分析 ===")

    # 使用util.py的load_dataset函数加载数据
    x_train, y_train = util.load_dataset(ds1_training_set_path, add_intercept=True)
    x_valid, y_valid = util.load_dataset(ds1_valid_set_path, add_intercept=True)

    # 训练逻辑回归（使用p01b中的实现）
    print("训练逻辑回归模型...")
    log_reg = LogisticRegression()
    log_reg.fit(x_train, y_train)

    # 训练GDA（使用p01e中的实现，需要去掉截距项）
    print("训练GDA模型...")
    x_train_no_intercept, _ = util.load_dataset(ds1_training_set_path, add_intercept=False)
    gda = GDA()
    gda.fit(x_train_no_intercept, y_train)

    # 计算准确率
    log_reg_acc = np.mean(log_reg.predict(x_valid) == y_valid)
    x_valid_no_intercept, _ = util.load_dataset(ds1_valid_set_path, add_intercept=False)
    gda_acc = np.mean(gda.predict(x_valid_no_intercept) == y_valid)

    print(f"数据集1 - 逻辑回归准确率: {log_reg_acc:.4f}")
    print(f"数据集1 - GDA准确率: {gda_acc:.4f}")

    # 使用扩展的plot函数绘制决策边界比较
    plot_comparison(
        x_valid_no_intercept, y_valid,
        log_reg.theta, gda.theta,
        'Dataset 1 - Decision Boundary Comparison',
        'output/Dataset1_comparison.png'
    )

    # ===== 数据集2分析 =====
    print("\n=== 数据集2分析 ===")

    # 使用util.py的load_dataset函数加载数据
    x_train, y_train = util.load_dataset(ds2_training_set_path, add_intercept=True)
    x_valid, y_valid = util.load_dataset(ds2_valid_set_path, add_intercept=True)

    # 训练逻辑回归（使用p01b中的实现）
    print("训练逻辑回归模型...")
    log_reg = LogisticRegression()
    log_reg.fit(x_train, y_train)

    # 训练GDA（使用p01e中的实现，需要去掉截距项）
    print("训练GDA模型...")
    x_train_no_intercept, _ = util.load_dataset(ds2_training_set_path, add_intercept=False)
    gda = GDA()
    gda.fit(x_train_no_intercept, y_train)

    # 计算准确率
    log_reg_acc_2 = np.mean(log_reg.predict(x_valid) == y_valid)
    x_valid_no_intercept, _ = util.load_dataset(ds2_valid_set_path, add_intercept=False)
    gda_acc_2 = np.mean(gda.predict(x_valid_no_intercept) == y_valid)

    print(f"数据集2 - 逻辑回归准确率: {log_reg_acc_2:.4f}")
    print(f"数据集2 - GDA准确率: {gda_acc_2:.4f}")

    # 使用扩展的plot函数绘制决策边界比较
    plot_comparison(
        x_valid_no_intercept, y_valid,
        log_reg.theta, gda.theta,
        'Dataset 2 - Decision Boundary Comparison',
        'output/Dataset2_comparison.png'
    )

    # ===== 结果分析 =====
    print("\n=== 结果分析 ===")
    print(f"数据集1 - 逻辑回归准确率: {log_reg_acc:.4f}, GDA准确率: {gda_acc:.4f}")
    print(f"数据集2 - 逻辑回归准确率: {log_reg_acc_2:.4f}, GDA准确率: {gda_acc_2:.4f}")

    print("\n=== 问题回答 ===")

    if gda_acc < log_reg_acc:
        print(f"\n答案：在数据集1上，GDA表现比逻辑回归差。")
        print(f"数据集1 - 逻辑回归准确率: {log_reg_acc:.4f}, GDA准确率: {gda_acc:.4f}")
        print("\n原因分析：")
        print("1. GDA假设数据服从高斯分布，但数据集1的数据分布可能不符合这个假设")
        print("2. 数据集1的特征分布可能不是正态分布，导致GDA的模型假设不成立")
        print("3. 逻辑回归没有对数据分布做特定假设，因此更灵活")
        print("4. 从notebook中可以看到，数据集1的x2特征都是非负的，这表明数据分布不是标准正态分布")
    else:
        print(f"\n在数据集1上，GDA表现与逻辑回归相当或更好。")

    if gda_acc_2 < log_reg_acc_2:
        print(f"\n在数据集2上，GDA表现比逻辑回归差。")
        print(f"数据集2 - 逻辑回归准确率: {log_reg_acc_2:.4f}, GDA准确率: {gda_acc_2:.4f}")
    else:
        print(f"\n在数据集2上，GDA表现与逻辑回归相当或更好。")
        print("这可能是因为数据集2的数据分布更接近高斯分布，符合GDA的假设。")


if __name__ == "__main__":
    main()